package xormplus

import (
	"database/sql/driver"
	"errors"
	"gitee.com/hongzhaomin/hzm-common-go/concurrent"
	"sync/atomic"
	"xorm.io/xorm"
)

type Transaction interface {
	driver.Tx
	Begin() error
	DoWithErr(fn func(session *xorm.Session) error) error
	Do(fn func(session *xorm.Session) error)
}

func NewTransaction(dataSourceId string) Transaction {
	return &transaction{
		dataSourceId: dataSourceId,
	}
}

type transaction struct {
	dataSourceId string
}

func (tx *transaction) Begin() error {
	engine := GetDataSourceById(tx.dataSourceId)
	session := engine.NewSession()
	// 存入goroutine local
	setGoSessionMap(engine, session)
	return session.Begin()
}

func (tx *transaction) Commit() error {
	session := tx.getSessionFromGoLocal()
	defer getSessionGoroutineLocal().Remove()
	return session.Commit()
}

func (tx *transaction) Rollback() error {
	session := tx.getSessionFromGoLocal()
	defer getSessionGoroutineLocal().Remove()
	return session.Rollback()
}

func (tx *transaction) getSessionFromGoLocal() *xorm.Session {
	engine := GetDataSourceById(tx.dataSourceId)
	return getGoSessionByEngine(engine)
}

func (tx *transaction) DoWithErr(fn func(session *xorm.Session) error) error {
	if err := tx.Begin(); err != nil {
		return err
	}

	return func() (e error) {
		defer func() {
			recovered := recover()
			if recovered != nil {
				var ok bool
				if e, ok = recovered.(error); !ok {
					e = errors.New("未知异常")
				}
				if err := tx.Rollback(); err != nil {
					logger().Error("== Transactional ==> 事务回滚异常: %v", err)
				}
			}
		}()
		if bizErr := fn(tx.getSessionFromGoLocal()); bizErr != nil {
			logger().Warn("== Transactional ==> 事务触发错误，进行回滚: %v", bizErr)
			if err := tx.Rollback(); err != nil {
				logger().Error("== Transactional ==> 事务回滚异常: %v", err)
				return err
			}
			return bizErr
		}
		return tx.Commit()
	}()
}

func (tx *transaction) Do(fn func(session *xorm.Session) error) {
	err := tx.DoWithErr(fn)
	if err != nil {
		logger().Error("== xormplus.transactional ==> 事务处理异常: %v", err)
		panic(err)
	}
}

var sessionGoroutineLocal atomic.Pointer[concurrent.GoroutineLocal[map[*xorm.Engine]*xorm.Session]]

func init() {
	sessionGoroutineLocal.Store(concurrent.NewGoroutineLocal[map[*xorm.Engine]*xorm.Session]())
}

func getSessionGoroutineLocal() *concurrent.GoroutineLocal[map[*xorm.Engine]*xorm.Session] {
	return sessionGoroutineLocal.Load()
}

func setGoSessionMap(engine *xorm.Engine, session *xorm.Session) {
	m := getSessionGoroutineLocal().Get()
	if m == nil {
		m = map[*xorm.Engine]*xorm.Session{
			engine: session,
		}
	} else {
		m[engine] = session
	}
	getSessionGoroutineLocal().Set(m)
}

func getGoSessionByEngine(engine *xorm.Engine) *xorm.Session {
	m := getSessionGoroutineLocal().Get()
	if m == nil {
		return nil
	}
	return m[engine]
}
